import os
import cv2
import time
import utils
import random
from tqdm import tqdm


def image_smoothing(BigImg, SmallImg, Infos, offsetx):
    if len(Infos) == 1:
        left, up, right, down = Infos[0][0] - offsetx, Infos[0][1], SmallImg.shape[1] - Infos[0][2] - offsetx, \
                                SmallImg.shape[0] - Infos[0][3]
        if left < 0 or right < 0:
            print('ddd')
    else:
        return SmallImg
        # xmin, ymin, xmax, ymax = SmallImg.shape[1], SmallImg.shape[0], 0, 0
        # for info in Infos:
        #     if info[0] < xmin:
        #         xmin = info[0]
        #     if info[1] < ymin:
        #         ymin = info[1]
        #     if info[2] > xmax:
        #         xmax = info[2]
        #     if info[3] > ymax:
        #         ymax = info[3]
        # left, up, right, down = xmin - offsetx, ymin, SmallImg.shape[1] - xmax - offsetx, SmallImg.shape[0] - ymax
    ran = 2
    for i in range(ran):
        y1 = int(i * up / ran)
        y2 = int(SmallImg.shape[0] - i * down / ran)
        x1 = int(i * left / ran)
        x2 = int(SmallImg.shape[1] - i * right / ran)
        BigImg[y1: y2, x1: x2] = cv2.addWeighted(BigImg[y1: y2, x1: x2], 1 - (i + 1) / ran, SmallImg[y1: y2, x1: x2], (i + 1) / ran,
                                                 0)
    return BigImg


SrcPath = r'/media/fang/TOSHIBA EXT/Temp/720*360'
DstPath = r'/media/fang/TOSHIBA EXT/Temp/M/720X360'
BackGroundPicPath = r'/media/fang/TOSHIBA EXT/BackGroundpic/road'
DstW = 720
DstH = 360
MixNum = 3

MaxW = int(DstW / MixNum)
MaxH = DstH
CutPicW = MaxW
CutPicH = MaxH
DstLocation = -1
NameNum = 0
names = []
BackGroundPics = []
for root, dirs, files in os.walk(BackGroundPicPath):
    for file in files:
        BackGroundPics.append(root + '/' + file)
for root, dirs, files in os.walk(SrcPath):
    for file in files:
        if file.split(".")[-1] == 'jpg':
            names.append(root + '/' + file)

random.shuffle(names)
random.shuffle(names)
for file in tqdm(names):
    # t1 = time.time()
    img = cv2.imread(file)
    locations = utils.ReadXml(file[:-3] + 'xml')
    for Info in locations:

        DstLocation += 1
        DstLocation = DstLocation % MixNum
        if DstLocation == 0:
            MixNum = random.randint(1, 4)
            MaxW = int(DstW / MixNum)
            NameNum += 1
            OutputLocations = [DstPath + '/M%06d.xml' % NameNum]
            OutImg = cv2.imread(random.choice(BackGroundPics))
            OutImg = cv2.resize(OutImg, (DstW, DstH))

        LocationW = Info[2] - Info[0]
        LocationH = Info[3] - Info[1]

        RandW = random.randint(0, min(MaxW - LocationW - 1, Info[0]))
        RandH = random.randint(0, min(MaxH - LocationH - 1, Info[1]))

        NewLocations = utils.CutInfos(Info[0] - RandW, Info[1] - RandH, Info[0] - RandW + CutPicW,
                                      Info[1] - RandH + CutPicH, 'NewLocation0', locations, CoverRate=0.2)

        for NewLocation in NewLocations[1:]:
            NewLocation[0] += MaxW * DstLocation
            NewLocation[2] += MaxW * DstLocation
            OutputLocations.append(NewLocation)

        imgNew = img[Info[1] - RandH:min(Info[1] - RandH + CutPicH, img.shape[0]),
                 Info[0] - RandW:min(Info[0] - RandW + CutPicW, img.shape[1])]
        OutImg[:imgNew.shape[0], MaxW * DstLocation: MaxW * DstLocation + imgNew.shape[1]] = imgNew
        # OutImg[:imgNew.shape[0], MaxW * DstLocation: MaxW * DstLocation + imgNew.shape[1]] = image_smoothing(
        #     OutImg[:imgNew.shape[0], MaxW * DstLocation: MaxW * DstLocation + imgNew.shape[1]], imgNew,
        #     OutputLocations[1:], MaxW * DstLocation)

        if DstLocation == MixNum - 1:
            cv2.imwrite(OutputLocations[0][:-3] + 'jpg', OutImg)
            utils.WriteXml(OutputLocations, DstW, DstH)
    # print(time.time() - t1)
    if file == names[-1] and DstLocation < MixNum:
        cv2.imwrite(OutputLocations[0][:-3] + 'jpg', OutImg)
        utils.WriteXml(OutputLocations, DstW, DstH)
